../../../../models/demos/wormhole/stable_diffusion/tests/test_transformer_2d_model.py